#!/usr/bin/env python
# encoding: utf-8
"""
untitled.py

Created by Orestis Markou on 2008-02-09.
Copyright (c) 2008 Orestis Markou. All rights reserved.
"""

import sys
import os
import unittest
import numpy
from numpy import array, zeros, ones, dot, argmax, max as amax
try:
    import cPickle as pickle
except ImportError:
    import pickle
    
class Picklable(object):
    @classmethod
    def read_model(cls, datafile):
        "read a model from a file"
        try:
            f =  open(datafile, 'rb')
            model = pickle.load(f)
        except:
            return None
        return model
    
    @classmethod        
    def write_model(cls, model,datafile):
        file =  open(datafile, 'wb')
        model = pickle.dump(model,file,pickle.HIGHEST_PROTOCOL)

class HMM(Picklable):
    """docstring for HMM"""
    def __init__(self, state_list, observation_list, A,B,pi):
        super(HMM, self).__init__()
        self.omega_X = state_list
        self.N= len(state_list)
        self.omega_O = observation_list
        self.M = len(observation_list)
        self.A = array(A)
        self.B = array(B)
        self.pi = array(pi)
        
        self.assertShapes()
    
    def assertShapes(self):
        assert(self.A.shape==(self.N,self.N))
        assert(self.B.shape==(self.M,self.N))
        assert(self.pi.shape==(self.N,))
    
    def __str__(self):
        return 'new_hmm with %d states and %d observations'%(self.N,self.M)

    def print_me(self):
        print 'A'*20
        print self.A
        print 'B'*20
        print self.B
        print 'pi'*10
        print self.pi
        print '-'*20
        
    def obs_indices(self, obs):
        """docstring for obs_indices"""
        ind = [self.omega_O.index(o) for o in obs]
        return ind
        
    def forward(self, observation_list):
        """docstring for forward"""
        pi = self.pi
        A = self.A
        B = self.B
        T = len(observation_list)
        obs = self.obs_indices(observation_list)
        a_ = zeros((T, self.N))
        a_[0] = pi*B[obs[0],:]
        
        for t in range(1,T):
            a_[t,:] = dot(a_[t-1,:],A)*B[obs[t],:]
        prob_ = sum(a_[T-1,:])
        return prob_
    
    def viterbi(self, observation_list):
        pi = self.pi
        A = self.A
        B = self.B
        T = len(observation_list)
        obs = self.obs_indices(observation_list)
        d_ = zeros((T,self.N))
        psi_ = zeros((T,self.N))
        #32a
        d_[0,:] = pi*B[obs[0],:]
        #33a,b
        for t in range(1,T):
            max_daij=amax(d_[t-1,:]*A.T,1)
            psi_[t,:] = argmax(d_[t-1,:]*A.T,1)
            d_[t,:] = max_daij*B[obs[t],:]
        #34a,b
        # P_star_ = max(d_[T-1,:])
        Q_T_star_ = argmax(d_[T-1,:])
        #35
        res = []
        for t in reversed(range(0,T)):
            res.append(self.omega_X[Q_T_star_])            
            Q_T_star_ = int(psi_[t][Q_T_star_])
        
        res.reverse()
        return res

    def _viterbi(self, observation_list):
        pi = self.pi
        A = self.A
        B = self.B
        T = len(observation_list)
        obs = self.obs_indices(observation_list)
        d = zeros((T,self.N))
        psi = zeros((T,self.N)) 
        #32a,b
        for i in range(self.N):
            d[0][i] = pi[i]*B[obs[0]][i]
            psi[0][i]=0
        
        d_ = zeros((T,self.N))
        psi_ = zeros((T,self.N))
        
        d_[0,:] = pi*B[obs[0],:]
        assert((d==d_).all())
        assert((psi==psi_).all())
        
        #33a,b
        for t in range(1,T):
            for j in range(self.N):
                max_daij = 0.0
                
                for i in range(self.N):
                    tmp_daij=d[t-1][i]*A[i][j]
                    if tmp_daij>=max_daij:
                        max_daij=tmp_daij
                        psi[t][j] = int(i)
                d[t][j] = max_daij*B[obs[t]][j]
        
        
        for t in range(1,T):
            max_daij=amax(d_[t-1,:]*A.T,1)
            psi_[t,:] = argmax(d_[t-1,:]*A.T,1)
            d_[t,:] = max_daij*B[obs[t],:]
        
        assert((d==d_).all())
        assert((psi==psi_).all())
        
        
        #34a,b
        P_star = 0
        Q_T_star = -1
        for i in range(self.N):
            tmp_P_star = d[T-1][i]
            if tmp_P_star>P_star:
                P_star = tmp_P_star
                Q_T_star = int(i)
                
        
        P_star_ = max(d_[T-1,:])
        Q_T_star_ = argmax(d_[T-1,:])
        
        assert(P_star_==P_star)
        assert(Q_T_star_==Q_T_star)
        # print Q_T_star_, Q_T_star
        #35
        res = []
        for t in reversed(range(0,T)):
            res.append(self.omega_X[Q_T_star])            
            Q_T_star = int(psi[t][Q_T_star])        
        res.reverse()

        res_ = []
        for t in reversed(range(0,T)):

            res_.append(self.omega_X[Q_T_star_])            
            Q_T_star_ = int(psi_[t][Q_T_star_])
            
        res_.reverse()
        assert(res_==res)

        return res
        
    
        
            
    def _forward(self, observation_list):
        """docstring for forward"""
        pi = self.pi
        A = self.A
        B = self.B
        T = len(observation_list)
        obs = self.obs_indices(observation_list)
        a = zeros((T, self.N))
        # a1 = zeros(self.N)
        for i in range(self.N):
            a[0][i] = pi[i]*B[obs[0]][i]
        # print 'a',a
        a_ = zeros((T, self.N))
        a_[0] = pi*B[obs[0],:]
        # print 'a_',a_
        assert((a==a_).all())
        
        for t in range(1,T):
            for j in range(self.N):
                for i in range(self.N):
                    a[t][j] += a[t-1][i]*A[i][j]
                # print 'a[%d][%d]'%(t,j),a[t][j]
                a[t][j] *= B[obs[t]][j]
                
        # for t in range(1,T):
        #     for j in range(self.N):
        #         a_[t][j] = sum(a_[t-1,:]*A[:,j])
        #         a_[t][j] *= B[obs[t]][j]
        # 
        # assert((a==a_).all())
        
        # for t in range(1,T):
        #     for j in range(self.N):
        #         a_[t][j] = sum(a_[t-1,:]*A[:,j])
        #     a_[t,:] *= B[obs[t],:]    

        # for t in range(1,T):
        #     a_[t,:] = dot(a_[t-1,:],A)
        #     a_[t,:] *= B[obs[t],:]
            
        # assert((a==a_).all())

        for t in range(1,T):
            a_[t,:] = dot(a_[t-1,:],A)*B[obs[t],:]
        assert((a==a_).all())
        
        prob = 0
        for i in range(self.N):
            prob += a[T-1][i]
        
        prob_ = sum(a_[T-1,:])
        assert(prob_==prob)
        return prob_
        
    


class HMMTests1(unittest.TestCase):
    def setUp(self):
        test = HMM(['a','b'],['s1','s2','s3'],
                   array([[.3,.7],[.5,.5]]),
                   array([[.5,0],[.5,.5],[0,.5]]),
                   array([.9,.1]))
        self.hmm = test
        # test.print_me()
    
    def test_indices(self):
        """docstring for test_indices"""
        ind = self.hmm.obs_indices(['s1','s2','s2','s3','s1','s2'])
        self.assertEqual(ind,[0,1,1,2,0,1])
        try:
            ind = self.hmm.obs_indices(['s1','FAIL'])
            self.fail('should got an exception')
        except:
            pass
    
    def test_forward(self):
        """docstring for test_forward"""
        self.hmm._forward(['s1','s2','s3'])
        self.assertEqual(self.hmm.forward(['s1','s2','s3']),0.063)
        
class HMMTestLeeds(unittest.TestCase):
    """docstring for HMMTestLeeds"""
    def setUp(self):
        test = HMM(['s','c', 'r'],['dry','dryish','damp','soggy'],
                   array([  [.5,   .25,  .25],
                            [.375, .125, .375],
                            [.125, .675, .375]
                            ]),
                   array([  [.6 , .25, .05],
                            [.2 , .25, .1],
                            [.15, .25, .35],
                            [.05, .25, .50]
                            ]),
                   array([.63,.17, .20]))
        self.hmm  = test
    
    def test_forward(self):
        """docstring for test_forward"""
        self.hmm._forward(['dry','damp','soggy'])
        self.assertAlmostEqual(self.hmm.forward(['dry','damp','soggy']),0.027386915)
        self.assertAlmostEqual(self.hmm.forward(
        'dry,damp,soggy,dry,dry,dry,dry,soggy,dry,dry'.split(',')),6.0686284E-6)
    
    def test_viterbi(self):
        self.hmm._viterbi(['dry','damp','soggy'])
        self.hmm._viterbi('dry,damp,soggy,dryish,dryish,damp'.split(','))
        self.assertEqual(self.hmm.viterbi(['dry','damp','soggy']),['s','r','r'])
        self.assertEqual(self.hmm.viterbi('dry,damp,soggy,dryish,dryish,damp'.split(',')),
            's,r,r,c,s,r'.split(','))



if __name__ == '__main__':
    unittest.main()